# K线图样式 csv+mplfinance https://blog.csdn.net/qq_41437512/article/details/105319421
# tushare+mplfinance——最小案例 https://blog.csdn.net/ooobenooo/article/details/107754092
# datatime——时间与字符串处理 https://www.cnblogs.com/chenqionghe/p/12235277.html
# mplfinance官网文档优秀翻译 https://blog.csdn.net/wuwei_201/article/details/105781844
# dataframe基础操作 https://www.jianshu.com/p/4ff1d2b23ab3
from datetime import datetime, timedelta, date

# import matplotlib as mpl
import mplfinance as mpf
import pandas as pd
import tushare as ts
import numpy as np
# import threading as win

api = None
data = None
data_insert = None
sCode = None
g_start_date = None
g_end_date = None
times = None
tk = {}
keyInfo = {}
bug_fix_1 = 0  # 修复重复输入收盘价时旧的预测收盘价不被删除的BUG
"""
固定的股票代码和时间段
sCode = '300059.SZ'
g_start_date = '20200201'
g_end_date = '20200731'
times = str(datetime.strptime(g_start_date, '%Y%m%d'))[:-9] + \
    " ~ " + str(datetime.strptime(g_end_date, '%Y%m%d'))[:-9]
"""


def get_sCode():
    global sCode
    sCode = input("请按格式输入股票代码：%s\n" % ('300059.SZ'))
    if len(sCode) != 9:
        print("您输入了不正确的股票代码 >>> %s\n已修改股票代码为默认值 >>> 300059.SZ" % sCode)
        sCode = '300059.SZ'


def set_time(next=False):
    global g_start_date
    global g_end_date
    global times

    g_start_date = (datetime.now() - timedelta(270)).strftime('%Y%m%d')
    g_end_date = next_workday().strftime('%Y%m%d') \
        if next else datetime.now().strftime('%Y%m%d')
    times = str(datetime.strptime(g_start_date, '%Y%m%d'))[:-9] + \
        " ~ " + str(datetime.strptime(g_end_date, '%Y%m%d'))[:-9]


def init_tushare():
    global api
    ts.set_token('dd8ab61a5210d4fa6ed707fd36c2db22986e7d3909e28b71e58f2f99')
    api = ts.pro_api()


def get_data():
    global data
    global sCode
    global times
    global g_start_date
    global g_end_date
    set_time(next=False)
    # Columns: [ts_code, trade_date, open, high, low, close, pre_close, change, pct_chg, vol, amount]
    df = api.query(api_name='daily',
                   ts_code=sCode,
                   start_date=g_start_date,
                   end_date=g_end_date)[::-1]
    if df.empty:
        print("\n\nERR!!! 股票：%s 在 %s 的数据查询失败，请仔细检查股票代码\n\n" % (sCode, times))
        raise Exception("股票：%s 在 %s 的数据查询失败，请仔细检查股票代码" % (sCode, times))
    else:
        print("股票：%s 在 %s 的数据 已查得\n数据共有 %d 行\n" %
              (sCode, times, df.size / 11))
    df['date'] = df['trade_date'].apply(
        lambda x: datetime.strptime(x, '%Y%m%d'))
    data = df.loc[:, ['date', 'open', 'close', 'high', 'low', 'vol']]
    data = data.rename(columns={"vol": "volume"})
    data.set_index('date', inplace=True)
    data.index.name = "Date"
    data.index = pd.DatetimeIndex(data.index)
    get_keyInfo(data)
    # data.shape // 可打印 data 的 shape
    # print(df)
    # print(data)


def draw(castmode=False):
    # show_nontrading:是否显示非交易日，默认False
    global data
    global sCode
    global times

    set_time(next=castmode)

    # print(mpl.rcParams)
    # mpl.rcParams["font.family"] = 'Arial Unicode MS' # 据说这是mac上设置中文 然并卵
    s = mpf.make_mpf_style(base_mpf_style='nightclouds', marketcolors=mpf.make_marketcolors(up='r', down='g', edge='i', wick='i'), mavcolors=[
        '#ffffff', '#f7d652', '#eb60f9', '#adadad', '#0647ef'])
    # mpl.rcParams['lines.linewidth'] = 20
    # mpl.rcParams['figure.dpi'] = 2000
    # mpl.rcParams['savefig.dpi'] = 2000

    tableName = "%s\n%s" % (sCode, times)
    # figpath = ('/Users/wangweijie/Desktop/股票预测走势绘制/Forcast:%s:%s' % (sCode, times) + '.png') \
    #     if castmode else'/Users/wangweijie/Desktop/股票预测走势绘制/%s:%s' % (sCode, times) + '.png'
    config = dict(
        type='candle', style=s, mav=(
            5, 10, 20, 60, 120), figratio=(5, 2), title=tableName, figscale=1.4, datetime_format='%m/%d', xrotation=0
    )

    print("股票：%s 在 %s 的数据 已读取 \n现有数据共有 %d 行\n" % (sCode, times, data.size / 5))
    print(data.head(3))
    print(data.tail(3))
    # mpf.plot(data, **config, savefig=figpath)
    if castmode:
        mpf.plot(data, **config)
    else:
        print("\n\n是否希望看到 原始数据 的走势图")
        print("Yes | yes | Y | y >>> 是")
        print("其他内容（例如回车） >>> 不，我选择不看原始数据图")
        key = input("请根据提示输入:  ")
        if key == "Yes" or key == "yes" or key == "Y" or key == "y":
            print("将弹出原始数据的走势图")
            print("按q关闭图片窗口")
            print("在图片窗口被关闭前，将无法开始预测")
            mpf.plot(data, **config)
        else:
            print("跳过展示 原始数据 的走势图")


def next_workday():
    today = datetime.now()
    if (today.weekday() == 4):  # friday
        nextwork_day = date.today() + timedelta(days=3)
    elif (today.weekday() == 5):  # sat
        nextwork_day = date.today() + timedelta(days=2)
    else:
        nextwork_day = date.today() + timedelta(days=1)
    return nextwork_day


def get_tk(mode):
    global tk
    if mode == 'B':
        while True:
            try:
                see_keyInfo()
                print("\n已设置收盘价：%s" % tk['close'])
                print("请按提示仔细输入")
                tk['open'] = float(input("请输入：开盘价\n"))
                tk['high'] = float(input("请输入：当日最高价\n"))
                tk['low'] = float(input("请输入：当日最低价\n"))
            except ValueError:
                print("输入值未满足格式要求：自然数或小数")
                pass  # 如果输入不能转换成浮点数，就重新来触发输入
            else:
                print(tk)
                if not (tk['high'] >= tk['open'] and tk['high'] >= tk['close'] and tk['low'] <= tk['open'] and tk['low'] <= tk['close'] and tk['low'] <= tk['high']):
                    print("输入值未满足大小关系：最低价 <= (开盘价,收盘价) <= 最高价")
                    pass  # 大小关系不满足，则重新输入
                else:
                    print("成功设置：开盘价、收盘价、当日最高价、当日最低价 ")
                    break  # 一切顺利，完成输入，退出死循环
    elif mode == 'A':
        while True:
            try:
                see_keyInfo()
                print("\n输入收盘价即可更新均线")
                tk['close'] = float(input("请输入：收盘价\n"))
            except ValueError:
                print("输入值未满足格式要求：自然数或小数")
                pass  # 如果输入不能转换成浮点数，就重新来触发输入
            else:
                tk['open'] = tk['high'] = tk['low'] = tk['close']
                print("成功设置：收盘价")
                print(tk)
                break  # 一切顺利，完成输入，退出死循环
    return tk


def forecast_data():
    global data
    global data_insert
    global bug_fix_1
    tk = get_tk(mode='A')
    key = input(
        "收盘价已设置：%s\n输入 No | no | N | n 则不渲染图片\n输入其他值或直接回车即默认渲染图片\n" % (tk['close']))
    if not (key == "No" or key == "no" or key == "N" or key == "n"):
        data_insert = pd.DataFrame(
            index=[next_workday()],
            data={
                'open': [tk['open']],
                'close': [tk['close']],
                'high': [tk['high']],
                'low': [tk['low']],
                'volume': [0]
            })
        if bug_fix_1 == 0:
            bug_fix_1 = 1
        else:
            data = data[:-1]  # 删除最后一行
        data = data.append(data_insert, ignore_index=False)
        data.index = pd.DatetimeIndex(data.index)
        draw(castmode=True)
    while True:
        print("是否继续预测")
        print("No | no | N | n >>> 不，结束此次预测")
        print("else >>> 是")
        key = input()
        if not (key == "No" or key == "no" or key == "N" or key == "n"):
            tk = get_tk(mode='B')
            data_insert = pd.DataFrame(
                index=[next_workday()],
                data={
                    'open': [tk['open']],
                    'close': [tk['close']],
                    'high': [tk['high']],
                    'low': [tk['low']],
                    'volume': [0]
                })
            data = data[:-1]  # 删除最后一行
            data = data.append(data_insert, ignore_index=False)
            data.index = pd.DatetimeIndex(data.index)
            draw(castmode=True)
        else:
            print("本轮预测已结束")
            break


def get_keyInfo(data):
    # x=mav5
    # x=mav10
    # x=涨停 则mav5
    global keyInfo

    # df = data.iloc[-20:].loc[:, ['close']] # 创建新的 df 内含 data 的 倒数二十行数据 的 'close' 这一列
    # dl 是 data 最后的二十个收盘价数据
    dl = []
    for x in np.array(data.iloc[-20:].loc[:, ['close']]).tolist():
        dl.append(x[0])
    # 不妨将dl倒序一下
    dl = dl[::-1]
    sum = 0
    for i in range(0, 4):
        sum += dl[i]
    keyInfo['mav5'] = sum / 4.0

    sum = 0
    for i in range(0, 9):
        sum += dl[i]
    keyInfo['mav10'] = sum / 9.0

    sum = 0
    for i in range(0, 19):
        sum += dl[i]
    keyInfo['mav20'] = sum / 19.0


def see_keyInfo():
    global keyInfo
    print("\n预测均线：")
    print("5日均线：%.2f" % keyInfo['mav5'])
    print("10日均线：%.2f" % keyInfo['mav10'])
    print("20日均线：%.2f" % keyInfo['mav20'])


if __name__ == "__main__":
    get_sCode()
    set_time()  # 默认值 next 为 false，即，并不会设置为下一个交易日
    init_tushare()  # tushare 为了能读取数据，需要初始化
    get_data()  # 根据读取今天往前数270天的数据，约为180个交易日
    draw()  # 默认值 castmode 为 false，即，并不会绘制和或保存预测数据的图片
    while True:
        print("\n\n是否要输入新的预测值")
        print("Yes | yes | Y | y >>> 是")
        print("No | no | N | n >>> 不，我选择退出程序")
        key = input("请根据提示输入:  ")
        if key == "Yes" or key == "yes" or key == "Y" or key == "y":
            print("开始新的一次预测")
            print("在预测图渲染图片窗口弹出后，可以按q关闭图片窗口")
            print("在图片窗口被关闭前，将无法开始下一轮预测")
            forecast_data()  # 触发预测数据
        elif key == "No" or key == "no" or key == "N" or key == "n":
            print("退出预测 程序中止")
            break
        else:
            print("输入值不合法 <<< %s 请按提示重新输入" % key)
